// third party c includes
#include <mpi.h>
#include <math.h>
#include <time.h>
#include <float.h>

// third party c++ includes
#include <iostream>
#include <iomanip>
#include <sstream>
#include <assert.h>
#include <strings.h>

// local includes
#include "flat_minable.h"
#include "param.h"
#include "data.h"
#include "sparse_array.h"
#include "lbfgs.h"
#include "util.h"

inline double log_sum(double a, double b);

//////////////////
//
// FlatMinable
//
//////////////////
FlatMinable::FlatMinable(EventsPtr data, GPParameter *param)
: CRFMinable(data, param)
{}

FlatMinable::~FlatMinable()
{}

double FlatMinable::objective_and_gradient(
	boost::shared_array<double> alpha, double &objective, 
	boost::shared_array<double> gradient)
{
	if (m_debug)  std::cout << "-- Calculating objective and gradient..."; 
	time_t start_t = time(0);

	objective = 0;
    if (m_rank == 0)
    {
		// global gradient contributions
		double inv_ss = 0.0;
		if (m_param->sigma_sqr != 0.0)
			inv_ss = 1.0 / m_param->sigma_sqr;

		// add the constant terms to G
		for(size_t x=0; x < m_param->n; ++x)
			gradient[x] = m_emperical_count[x] + inv_ss*m_weights[x];

		// global prior contributions
		double theta_norm = 0.0;
		for(size_t x=0; x < m_param->n; ++x)
			theta_norm += m_weights[x]*m_weights[x];
		objective = 0.5 * inv_ss * theta_norm; // the prior
	}
	else
		bzero(gradient.get(), sizeof(double)*m_param->n);

	double acc=0;
	FlatEvents* flat_events = dynamic_cast<FlatEvents*>(m_data.get());
	boost::shared_array<double> potentials = flat_events->product(alpha);
	boost::shared_array<double> marginals
		(new double [flat_events->labels()*flat_events->events()]);

	// calculate objective and marginals
	int curr_index=0;
	for (int i=0; i<flat_events->events(); ++i)
	{
		// calc Z
		double log_z=-DBL_MAX;
		for (int l=0; l<flat_events->labels(); ++l)
			log_z = log_sum(log_z, potentials[curr_index+l]);
		// calc marginals
		int gold = flat_events->label(i);
		double potential=0;
		for (int l=0; l<flat_events->labels(); ++l)
		{
			potential = potentials[curr_index+l] - log_z;
			marginals[curr_index+l] = exp(potential);
			if (l == gold) objective -= potential;
		}
		curr_index += flat_events->labels();
	}

	// calculate gradient
	boost::shared_array<double> expectations = flat_events->transpose_product(marginals);
	for(size_t x=0; x < m_param->n; ++x)
		gradient[x] += expectations[x];

// ////////////////////////////////////
// std::cout << "Probs: " << std::endl;
// for (unsigned int i=0; i < 9; ++i)
//  std::cout << " " << marginals[i];
// std::cout << std::endl;
// ////////////////////////////////////
// //////////////////////////////////////
// std::cout << "Gradient: " << std::endl;
// for (unsigned int i=0; i < 106; ++i)
//  std::cout << " " << gradient[i];
// std::cout << std::endl;
// //////////////////////////////////////
// //////////////////////////////////////
// std::cout << "Weights: " << std::endl;
// for (unsigned int i=0; i < 106; ++i)
//  std::cout << " " << m_weights[i];
// std::cout << std::endl;
// //////////////////////////////////////
	
	// record some trace data
	time_t end_t = time(0);
	double elapsed_time = difftime(end_t,start_t);
	m_it_data.iteration_secs += (int) elapsed_time;
	m_it_data.total_secs += (int) elapsed_time;
	m_it_data.function_evals += 1;

	if (m_debug) 
		std::cout << " done in: " <<	elapsed_time << " secs" << std::endl;

	return acc;
}

void FlatMinable::_set_emperical_kernel_count( void )
{
	if (m_debug) std::cout << "  ->setting emperical counts" << std::endl;
	time_t start_t = time(0);

	// create a unit vector to sum features
	FlatEvents* flat_events = dynamic_cast<FlatEvents*>(m_data.get());
	int num_rows = flat_events->events()*flat_events->labels();
	boost::shared_array<double> unit_vector(new double[num_rows]);
	bzero(unit_vector.get(), sizeof(double)*num_rows);
	for(int x=0; x < flat_events->events(); ++x) 
		unit_vector[x*flat_events->labels()+flat_events->label(x)] = -1.0f;

	// a matrix product does the sum
	boost::shared_array<double> local_emperical_count=flat_events->transpose_product(unit_vector);

	if (m_rank == 0)
		m_emperical_count = boost::shared_array<double>(new double[m_param->n]);
	if (MPI_SUCCESS != MPI_Reduce(local_emperical_count.get(), 
		m_emperical_count.get(), m_param->n, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD))
	MPI_Abort(MPI_COMM_WORLD,1);

//	//////////////////////////////////////////////
//	std::cout << "Emperical Counts:" << std::endl;
//	for (unsigned int i=0; i < m_param->n && i < 1000; ++i)
//		std::cout << " " << m_emperical_count[i];
//	std::cout << std::endl;
//	//////////////////////////////////////////////

	time_t end_t = time(0);
	double elapsed_time = difftime(end_t,start_t);
	if (m_debug) 
		std::cout << "\n <-done in: " << elapsed_time << " secs" << std::endl;
}

///////////////////////////////////////////////////////////////////////////////
// log(exp(a) + exp(b))
///////////////////////////////////////////////////////////////////////////////
inline double
log_sum(double a, double b)
{
	if (a>b) return a + log(1+exp(b-a));
	else	 return b + log(1+exp(a-b));
}
